import os
import re
import csv
from datetime import datetime, timedelta
from pathlib import Path
"""
刚分红过的股票,一般不会马上再分红
为了应对一年分红多次的情况, 排除范围定为90天
但有时又想快速获取当天的分红公告,所以排队范围可定为275天
"""
exclude_days = 90

def collect_recent_csv_files() -> set[str]:
    """
    收集当前目录下所有 .csv 文件中创建日期在90天以内的文件名
    
    返回:
        符合条件的文件名集合
    """
    recent_files = set()
    current_dir = os.getcwd()
    
    for entry in os.scandir(current_dir):
        if entry.is_file() and entry.name.lower().endswith('.csv'):
            # 从文件名中提取日期
            date_match = re.search(r'(\d{4}-\d{2}-\d{2})', entry.name)
            if date_match:
                file_date_str = date_match.group(1)
                
                try:
                    # 转换日期字符串为 datetime 对象
                    file_date = datetime.strptime(file_date_str, '%Y-%m-%d').date()
                    
                    # 计算90天前的日期
                    ninety_days_ago = datetime.now().date() - timedelta(days=exclude_days)
                    
                    # 判断文件日期是否在90天内
                    if file_date >= ninety_days_ago:
                        recent_files.add(entry.name)
                        
                except ValueError:
                    print(f"警告: 无法解析文件 {entry.name} 中的日期格式")
    
    return recent_files

def extract_stock_ids_from_csv(file_path: str) -> set[str]:
    """
    从 CSV 文件的第一列提取股票 ID
    
    参数:
        file_path: CSV 文件路径
    
    返回:
        股票 ID 集合
    """
    stock_ids = set()
    
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            reader = csv.reader(file)
            # # 不跳过标题行 因为没有标题行
            # next(reader, None)
            
            for row in reader:
                if row:  # 确保行不为空
                    stock_ids.add(row[0].strip())  # 取第一列并去除空白字符
                    
    except Exception as e:
        print(f"错误: 无法处理文件 {file_path} - {e}")
        
    return stock_ids

def save_stock_ids_to_file(stock_ids: set[str], output_file: str = 'exclude_stock_ids.txt') -> None:
    """
    将股票 ID 集合保存到文本文件
    
    参数:
        stock_ids: 股票 ID 集合
        output_file: 输出文件名
    """
    try:
        with open(output_file, 'w', encoding='utf-8') as file:
            # 按字母顺序排序并写入文件，每行一个 ID
            for stock_id in sorted(stock_ids):
                file.write(f"{stock_id}\n")
                
        print(f"成功保存 {len(stock_ids)} 个股票 ID 到 {output_file}")
        
    except Exception as e:
        print(f"错误: 无法保存股票 ID 到文件 - {e}")

def main():
    # 步骤1: 收集符合条件的 CSV 文件
    recent_files = collect_recent_csv_files()
    
    if not recent_files:
        print("未找到符合条件的 CSV 文件")
        return
    
    print(f"找到 {len(recent_files)} 个符合条件的 CSV 文件")
    
    # 步骤2: 提取所有股票 ID
    all_stock_ids = set()
    
    for filename in recent_files:
        file_path = os.path.join(os.getcwd(), filename)
        stock_ids = extract_stock_ids_from_csv(file_path)
        
        if stock_ids:
            print(f"从 {filename} 中提取了 {len(stock_ids)} 个股票 ID")
            all_stock_ids.update(stock_ids)
        else:
            print(f"警告: {filename} 不包含有效的股票 ID")
    
    # 步骤3: 保存股票 ID 到文件
    if all_stock_ids:
        save_stock_ids_to_file(all_stock_ids)
    else:
        print("没有提取到任何股票 ID")

if __name__ == "__main__":
    main()    